import cv2 as cv
import numpy as np


def decode_box(x):
    """输出解码的numpy版本，用于onnx"""
    x = [x[..., :300], x[..., 300:1500], x[..., 1500:]]
    x = [x[0].reshape(-1, 5, 15, 20), x[1].reshape(-1, 5, 30, 40), x[2].reshape(-1, 5, 60, 80)]
    predictions = []
    def sigmoid(s): return 1 / (1 + np.exp(-s))
    for i in range(3):
        predict = x[i].transpose(0, 2, 3, 1)  # 批大小,高,宽,预测结果 N,H,W,P
        img_height, img_width = predict.shape[1], predict.shape[2]
        grid_x = np.expand_dims(np.arange(0, img_width), axis=0).repeat(img_height, 0)
        grid_y = np.expand_dims(np.arange(0, img_height), axis=0).repeat(img_width, 0).transpose(1, 0)
        predict[..., 0] = sigmoid(predict[..., 0]) + grid_x  # 横坐标
        predict[..., 1] = sigmoid(predict[..., 1]) + grid_y  # 纵坐标
        predict[..., 2] = 7 * np.exp(predict[..., 2])  # 宽度
        predict[..., 3] = 10 * np.exp(predict[..., 3])  # 高度
        predict[..., 4] = sigmoid(predict[..., 4])  # 置信度
        predictions.append(predict)
    return tuple(predictions)


def image_generate_loc(img, label):
    """在图像上显示结果
    :param img:网络输入的原图像
    :param label:网络经解码后的输出
    """
    stride = 32, 16, 8
    cnt = [0, 0, 0]
    predict = []
    for fm in range(3):
        predict.append(label[fm].reshape(-1, 5))
        cnt[fm] = predict[fm].shape[0]
    predict = np.vstack(predict)
    index = predict[:, 4].argmax()
    predict = predict[index]
    if predict[4] < 0.4:
        return img
    corner = predict[:4].copy()
    corner[0] = predict[0] - predict[2] / 2
    corner[1] = predict[1] - predict[3] / 2
    corner[2] = predict[0] + predict[2] / 2
    corner[3] = predict[1] + predict[3] / 2
    fm = 0
    while index > 0:
        index -= cnt[fm]
        fm += 1
    fm -= 1
    corner = np.round(corner * stride[fm]).astype('int32')
    img = cv.rectangle(img, (corner[0], corner[1]), (corner[2], corner[3]), color=(0, 0, 255), thickness=2)
    pred_conf = '{:.4f}'.format(predict[4])
    img = cv.putText(img, pred_conf, (corner[0], corner[1]), cv.FONT_HERSHEY_PLAIN, 1.2, (0, 0, 255), 1)
    return img


print('Start detect via camera.')
cap = cv.VideoCapture(0)
cap.set(cv.CAP_PROP_FRAME_WIDTH, 1280)
cap.set(cv.CAP_PROP_FRAME_HEIGHT, 480)
net = cv.dnn.readNetFromONNX('zhnnet.onnx')
while cap.isOpened():
    ret, image_origin = cap.read()
    if not ret:
        break
    imgL = image_origin[:, :640, :]
    imgR = image_origin[:, 640:, :]
    image = image_origin.copy()
    imgL_x = imgL.transpose(2, 0, 1) / 256
    imgR_x = imgR.transpose(2, 0, 1) / 256
    imgL_x = np.expand_dims(imgL_x, axis=0)
    imgR_x = np.expand_dims(imgR_x, axis=0)
    image = np.vstack([imgL_x, imgR_x])
    net.setInput(image)
    pred = net.forward()
    pred = decode_box(pred)
    imgLpred = pred[0][0], pred[1][0], pred[2][0]
    imgRpred = pred[0][1], pred[1][1], pred[2][1]
    imgL = image_generate_loc(imgL, imgLpred)
    imgR = image_generate_loc(imgR, imgRpred)
    image_origin[:, :640, :] = imgL
    image_origin[:, 640:, :] = imgR
    cv.imshow('test', image_origin)
    c = cv.waitKey(10)
    if c == ord('q'):
        break
    elif c == ord('s'):
        cv.imwrite('err.png', image)
cv.destroyAllWindows()
exit()
